package org.bdware.doip.core.codec.packet;

import com.google.gson.JsonObject;
import com.nimbusds.jose.jwk.JWK;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import org.apache.log4j.Logger;
import org.bdware.doip.core.crypto.GlobalCertifications;
import org.bdware.doip.core.doipMessage.DoipMessage;
import org.bdware.doip.core.doipMessage.HeaderParameter;
import org.bdware.doip.core.doipMessage.MessageCredential;
import org.bdware.doip.core.doipMessage.MessageEnvelope;
import org.bdware.doip.core.exception.MessageCodecException;
import org.bdware.doip.core.utils.DoipGson;
import org.bdware.irp.client.GlobalIrpClient;
import org.bdware.irp.crypto.CertUtils;
import org.bdware.irp.exception.IrpClientException;
import org.bdware.irp.stateinfo.UserStateInfo;

import java.net.InetSocketAddress;
import java.util.ArrayList;

public class PacketMessageCodecImpl implements PacketMessageCodec{
    Logger logger = Logger.getLogger(PacketMessageCodecImpl.class);

    @Override
    public byte[] EnvelopeToBytes(MessageEnvelope envelope) throws MessageCodecException {
        if(envelope.contentLength != envelope.content.length){
            throw new MessageCodecException("unequal content length");
        }
        ByteBuf bf = Unpooled.directBuffer();
        bf.writeByte(envelope.majorVersion);
        bf.writeByte(envelope.minVersion);
        bf.writeShort(envelope.getFlag());
        bf.writeInt(envelope.reserved);
        bf.writeInt(envelope.requestId);
        bf.writeInt(envelope.sequenceNumber);
        bf.writeInt(envelope.totalNumber);
        bf.writeInt(envelope.contentLength);
        bf.writeBytes(envelope.content);
        byte[] encodedEnvelope = new byte[bf.readableBytes()];
        bf.readBytes(encodedEnvelope);
        bf.release();
        return encodedEnvelope;
    }

    @Override
    public MessageEnvelope BytesToEnvelope(byte[] envelopeBytes) throws MessageCodecException {

        ByteBuf bf = Unpooled.directBuffer();
        bf.writeBytes(envelopeBytes);
        MessageEnvelope envelope = new MessageEnvelope();
        envelope.majorVersion = bf.readByte();
        envelope.minVersion = bf.readByte();
        envelope.setFlag(bf.readShort());
        envelope.reserved = bf.readInt();
        envelope.requestId = bf.readInt();
        envelope.sequenceNumber = bf.readInt();
        envelope.totalNumber = bf.readInt();
        if(envelope.sequenceNumber > envelope.totalNumber)
            throw new MessageCodecException("invalid sequence number: " + envelope.sequenceNumber + ", total: " + envelope.totalNumber);
        envelope.contentLength = bf.readInt();
        if(envelope.contentLength != bf.readableBytes()){
            throw new MessageCodecException("unequal content length: " + envelope.contentLength + ":" +bf.readableBytes());
        }
        envelope.content = new byte[envelope.contentLength];
        bf.readBytes(envelope.content);
        bf.release();
        return envelope;

    }

    @Override
    public DoipMessage EnvelopesToMessage(ArrayList<MessageEnvelope> envelopList) throws MessageCodecException {
        if(envelopList.size() != envelopList.get(0).totalNumber)
            throw new MessageCodecException("invalid envelop number: list size: " + envelopList.size() + " totalNumber: " + envelopList.get(0).totalNumber);
        ByteBuf bf = Unpooled.directBuffer();
        for (MessageEnvelope en:envelopList) {
            bf.writeBytes(en.content);
        }
        byte[] msgBytes = new byte[bf.readableBytes()];
        bf.readBytes(msgBytes);
        DoipMessage msg;
        if(envelopList.get(0).isEncrypted()){
            try {
                logger.debug("decrypted msg data length: " + msgBytes.length);
                msgBytes = CertUtils.decrypt(msgBytes,GlobalCertifications.getGlobalJWK());
                logger.debug("raw msg data length: " + msgBytes.length);
                msg = BytesToMessage(msgBytes,envelopList.get(0).requestId,envelopList.get(0).getSender());
                msg.header.setIsEncrypted(true);
            } catch (Exception e) {
                logger.debug("decryption data failed.");
                throw  new MessageCodecException("decryption data failed.");
            }
        }else{
            msg = BytesToMessage(msgBytes,envelopList.get(0).requestId,envelopList.get(0).getSender());
        }
        bf.release();
        return msg;
    }

    @Override
    public ArrayList<MessageEnvelope> MessageToEnvelopes(DoipMessage msg) throws MessageCodecException {
        ByteBuf bf = Unpooled.directBuffer();
        byte[] msgBytes = MessageToBytes(msg);
        if(msg.header.isEncrypted() && msg.getRecipientID()!= null){
            try {
                UserStateInfo recipientInfo = new UserStateInfo(GlobalIrpClient.getGlobalClient().resolve(msg.getRecipientID()));
                logger.debug("recipient pk: " + recipientInfo.getPubkey());
                JWK recipientPK = JWK.parse(recipientInfo.getPubkey());
                logger.debug("raw msg data length: " + msgBytes.length);
                msgBytes = CertUtils.encrypt(msgBytes,recipientPK);
                logger.debug("encrypted msg data length: " + msgBytes.length);
            } catch (IrpClientException ie) {
                logger.warn("resolution recipient's pk failed, will not do encryption");
            }catch (Exception e) {
                logger.debug("encryption data error.");
                e.printStackTrace();
            }
        }
        bf.writeBytes(msgBytes);
        //buffer to envList
        ArrayList<MessageEnvelope> envelopes = new ArrayList<>();
        int encodedLength = bf.readableBytes();
        int totalNumber = encodedLength/msg.getMTU() +1;

        while(bf.isReadable()){
            MessageEnvelope env = new MessageEnvelope();
            if(bf.readableBytes() < msg.getMTU()){
                env.content = new byte[bf.readableBytes()];
                bf.readBytes(env.content);
                env.setTruncated(false);
            }else{
                env.content = new byte[msg.getMTU()];
                bf.readBytes(env.content);
                env.setTruncated(true);
            }
            env.setTruncated(totalNumber > 1);
            env.setEncrypted(msg.header.isEncrypted());
            env.sequenceNumber = envelopes.size();
            env.totalNumber = totalNumber;
            env.requestId = msg.requestID;
            env.contentLength = env.content.length;
            env.setSender(msg.getSender());
            envelopes.add(env);
        }
        bf.release();
        return envelopes;
    }

    @Override
    public byte[] MessageToBytes(DoipMessage msg) throws MessageCodecException {
        ByteBuf bf = Unpooled.directBuffer();
        msg.header.parameterLength = msg.header.parameters.length();
        msg.header.bodyLength = msg.body.getLength();

        //encode header
        bf.writeInt(msg.header.getFlag());
        bf.writeInt(msg.header.parameterLength);
        bf.writeInt(msg.header.bodyLength);
        if(msg.header.parameterLength != 0 && msg.header.parameterLength != msg.header.parameters.length())
            throw new MessageCodecException("invalid parameter length: " + msg.header.parameterLength);
        bf.writeBytes(msg.header.parameters.toByteArray());

        //encode body
        if(msg.header.bodyLength != 0 && msg.header.bodyLength != msg.body.getLength())
            throw new MessageCodecException("invalid body length: " + msg.header.parameterLength);
        bf.writeBytes(msg.body.getEncodedData());

        //encode credential
        if(msg.credential != null){
            bf.writeInt(msg.credential.attributeLength());
            bf.writeBytes(DoipGson.getDoipGson().toJson(msg.credential.attributes).getBytes());
            bf.writeInt(msg.credential.sigLength());
            bf.writeBytes(msg.credential.getSignature());
        }

        byte[] msgBytes = new byte[bf.readableBytes()];
        bf.readBytes(msgBytes);
        bf.release();
        return msgBytes;
    }

    @Override
    public DoipMessage BytesToMessage(byte[] msgBytes,int requestId, InetSocketAddress sender) throws MessageCodecException {
        ByteBuf bf = Unpooled.directBuffer();
        bf.writeBytes(msgBytes);
        //decode header
        DoipMessage msg = new DoipMessage("","");
        msg.requestID = requestId;
        msg.setSender(sender);
        msg.header.setFlag(bf.readInt());
        msg.header.parameterLength = bf.readInt();
        msg.header.bodyLength = bf.readInt();
        byte[] parameters = new byte[msg.header.parameterLength];
        bf.readBytes(parameters);
        msg.header.parameters = DoipGson.getDoipGson().fromJson(new String(parameters), HeaderParameter.class);
        logger.debug(bf.readableBytes());
        //decode body
        if (msg.header.bodyLength > 0) {
            if(bf.readableBytes() < msg.header.bodyLength) throw new MessageCodecException("invalid body length");
            msg.body.encodedData = new byte[msg.header.bodyLength];
            bf.readBytes(msg.body.encodedData);
        }
        logger.debug(bf.readableBytes());
        //decodeCredential
        if (bf.readableBytes() > 0){
            byte[] attrBytes = readDataArray(bf);
            logger.debug("string: " + new String(attrBytes));
            JsonObject attrJo = DoipGson.getDoipGson().fromJson(new String(attrBytes),JsonObject.class);
            byte[] signature = readDataArray(bf);
            msg.credential = new MessageCredential(attrJo);
            msg.credential.setSignature(signature);
        }
        bf.release();
        return msg;
    }

    private static byte[] readDataArray(ByteBuf din) throws MessageCodecException {
        int dataLen = din.readInt();
        if(dataLen < 0 || dataLen > din.readableBytes()) throw new MessageCodecException("invalid credential length");
        byte[] data = new byte[dataLen];
        din.readBytes(data);
        return data;
    }

}
